import frappe
from enum import Enum
from frappe.query_builder.builder import MariaDB, Postgres
from frappe.query_builder.utils import db_type_is
from frappe.query_builder.utils import ImportMapper
from frappe_dm.query_builder.builder import Dm

def patched_get_query_builder(type_of_db: str):
    from frappe.query_builder.utils import db_type_is
    """补丁后的 get_query_builder 函数"""
    db = frappe.query_builder.utils.db_type_is(type_of_db)
    picks = {
        db_type_is.MARIADB: MariaDB,
        db_type_is.POSTGRES: Postgres,
        db_type_is.DM: Dm
    }
    return picks[db]

def patch_query_builder():
    """应用查询构建器补丁"""
    import frappe.query_builder
    import frappe.query_builder.utils
    import frappe.query_builder.functions
    from frappe_dm.query_builder.custom import (
        DM_GROUP_CONCAT, DM_MATCH,
        DM_COMBINE_DATETIME, DM_DATE_FORMAT,
        DM_UNIX_TIMESTAMP
    )
    import sys
    
    # 添加达梦数据库类型到枚举
    original_values = [(e.name, e.value) for e in db_type_is]
    original_values.append(('DM', 'dm'))
    
    # 使用type动态创建新的枚举类
    NewDbTypeIs = Enum('db_type_is', dict(original_values))
    
    # 替换所有可能引用的地方
    frappe.query_builder.utils.db_type_is = NewDbTypeIs
    sys.modules['frappe.query_builder.utils'].db_type_is = NewDbTypeIs

    # 扩展 get_query_builder 函数支持达梦数据库
    print("扩展 get_query_builder 函数...")
    
    # 替换所有可能引用的地方
    frappe.query_builder.utils.get_query_builder = patched_get_query_builder
    frappe.query_builder.get_query_builder = patched_get_query_builder
    frappe.get_query_builder = patched_get_query_builder
    sys.modules['frappe.query_builder.utils'].get_query_builder = patched_get_query_builder
    sys.modules['frappe.query_builder'].get_query_builder = patched_get_query_builder
    sys.modules['frappe'].get_query_builder = patched_get_query_builder

    # 扩展数据库函数支持达梦数据库
    print("扩展数据库函数...")
    
    # 获取原始的 ImportMapper
    original_group_concat = frappe.query_builder.functions.GroupConcat
    original_match = frappe.query_builder.functions.Match
    original_combine_datetime = frappe.query_builder.functions.CombineDatetime
    original_date_format = frappe.query_builder.functions.DateFormat
    original_unix_timestamp = frappe.query_builder.functions.UnixTimestamp

    # 创建新的 ImportMapper 实例
    frappe.query_builder.functions.GroupConcat = ImportMapper({
        **original_group_concat.func_map,
        NewDbTypeIs.DM: DM_GROUP_CONCAT
    })

    frappe.query_builder.functions.Match = ImportMapper({
        **original_match.func_map,
        NewDbTypeIs.DM: DM_MATCH
    })

    frappe.query_builder.functions.CombineDatetime = ImportMapper({
        **original_combine_datetime.func_map,
        NewDbTypeIs.DM: DM_COMBINE_DATETIME
    })

    frappe.query_builder.functions.DateFormat = ImportMapper({
        **original_date_format.func_map,
        NewDbTypeIs.DM: DM_DATE_FORMAT
    })

    frappe.query_builder.functions.UnixTimestamp = ImportMapper({
        **original_unix_timestamp.func_map,
        NewDbTypeIs.DM: DM_UNIX_TIMESTAMP
    })

    # 修补 PseudoColumnMapper 的 get_sql 方法
    print("修补 PseudoColumnMapper...")
    original_get_sql = frappe.query_builder.utils.PseudoColumnMapper.get_sql
    
    def patched_get_sql(self, **kwargs):
        """修补后的 get_sql 方法，支持达梦数据库"""
        if frappe.db.db_type in ("postgres", "dm"):
            self.name = self.name.replace("`", '"')
        return self.name
    
    frappe.query_builder.utils.PseudoColumnMapper.get_sql = patched_get_sql

    print("查询构建器补丁应用完成") 